from typing import Dict, Callable
import os

import dashscope
from dashscope import Generation

from mm_story_agent.base import register_tool


@register_tool("qwen")
class QwenAgent(object):

    def __init__(self,
                 config: Dict):
        
        self.system_prompt = config.get("system_prompt")
        track_history = config.get("track_history", False)
        if self.system_prompt is None:
            self.history = []
        else:
            self.history = [
                {"role": "system", "content": self.system_prompt}
            ]
        self.track_history = track_history
        
        # 从配置或环境变量读取 API Key，并设置到 dashscope 模块级别
        self.api_key = config.get("api_key") or os.environ.get("DASHSCOPE_API_KEY")
        if self.api_key:
            # 设置到环境变量（供其他地方使用）
            os.environ["DASHSCOPE_API_KEY"] = self.api_key
            # 设置到 dashscope 模块级别（全局生效）
            dashscope.api_key = self.api_key
        else:
            # 如果没有设置 API Key，给出警告
            print("⚠️  警告: 未找到 DashScope API Key！")
            print("   请在配置文件中添加 'api_key' 字段，或设置环境变量 DASHSCOPE_API_KEY")
    
    def basic_success_check(self, response):
        if not response or not response.output or not response.output.text:
            print(response)
            return False
        else:
            return True
    
    def call(self,
             prompt: str,
             model_name: str = "qwen2-72b-instruct",
             top_p: float = 0.95,
             temperature: float = 1.0,
             seed: int = 1,
             max_length: int = 1024,
             max_try: int = 5,
             success_check_fn: Callable = None
             ):
        self.history.append({
            "role": "user",
            "content": prompt
        })
        success = False
        try_times = 0
        # 安全限制：DashScope 当前最大输出 token 为 6144
        safe_max_length = max(1, min(max_length, 6144))
        while try_times < max_try:
            # 构建调用参数
            call_kwargs = {
                "model": model_name,
                "messages": self.history,
                "top_p": top_p,
                "temperature": temperature,
                "seed": seed,
                "max_length": safe_max_length
            }
            # 如果有 api_key，传递给调用；否则使用模块级别的 key
            if self.api_key:
                call_kwargs["api_key"] = self.api_key
            
            try:
                response = Generation.call(**call_kwargs)
            except Exception as e:
                err_msg = str(e)
                if "Range of max_tokens" in err_msg and safe_max_length > 512:
                    safe_max_length = max(512, int(safe_max_length * 0.75))
                    print(f"⚠️  DashScope 限制 max_tokens，调整 max_length 为 {safe_max_length} 并重试")
                    try_times += 1
                    continue
                else:
                    raise

            if success_check_fn is None:
                success_check_fn = lambda x: True
            if self.basic_success_check(response) and success_check_fn(response.output.text):
                response = response.output.text
                self.history.append({
                    "role": "assistant",
                    "content": response
                })
                success = True
                break
            else:
                try_times += 1
        
        if not self.track_history:
            if self.system_prompt is not None:
                self.history = self.history[:1]
            else:
                self.history = []
        
        return response, success
   